import torch as pt
from torch.utils import data
import numpy as np
from python_ai.common.xcommon import *

M, N = 100, 10
BATCH_SIZE = 32
x = np.random.normal(0, 1, [M, N])
y = np.random.choice(np.arange(N, dtype=np.int64), size=M, p=np.full(N, 1 / N))
print_numpy_ndarray_info(x, 'x')
print_numpy_ndarray_info(y, 'y')

xt = pt.Tensor(x)
yt = pt.Tensor(y)
ds = data.TensorDataset(xt, yt)
dl = data.DataLoader(ds, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)

for i, (bx, by) in enumerate(dl):
    sep(i)
    bx = bx.numpy()
    by = by.numpy()
    print_numpy_ndarray_info(bx, 'bx')
    print_numpy_ndarray_info(by, 'by')
